import pandas as pd
from sklearn.linear_model import Lasso  # 导入Lasso回归
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error

df = pd.read_excel('F:/25MCM_C/C_data/USA_ROU.xlsx', engine='openpyxl')

df['coach_change'] = df.apply(lambda row: 1 if (row['NOC'] == 'USA' and row['YEAR'] >= 1981) or (
            row['NOC'] == 'ROU' and row['YEAR'] >= 1981) else 0, axis=1)

df['gold_count'] = df['GOLD']  # 只使用金牌数

print(f"Unique NOC values: {df['NOC'].unique()}")
#比赛结束前最后一天售后群发布无水印可视化结果+无标注代码【可直接提交】为了防止倒卖， 论文写作过程中遗留数个致命问题，无关代码，该问题解决方式仅在官网授权售后群答疑，盗卖方式购买资料不提供答疑，感谢理解 美赛资料助攻购买链接+说明https://docs.qq.com/doc/p/f3dc6bffbf4dab58dbdfd3e5e5de18a2ad974216
df_usa = df[df['NOC'] == 'USA']
df_rou = df[df['NOC'] == 'ROU']

print(f"Number of USA records: {df_usa.shape[0]}")
print(f"Number of Romania records: {df_rou.shape[0]}")
#比赛结束前最后一天售后群发布无水印可视化结果+无标注代码【可直接提交】为了防止倒卖， 论文写作过程中遗留数个致命问题，无关代码，该问题解决方式仅在官网授权售后群答疑，盗卖方式购买资料不提供答疑，感谢理解 美赛资料助攻购买链接+说明https://docs.qq.com/doc/p/f3dc6bffbf4dab58dbdfd3e5e5de18a2ad974216
if df_rou.shape[0] == 0:
    print("Warning: No data for Romania (NOC == 'ROU')")

def prepare_data(df):
    X = pd.get_dummies(df[['coach_change', 'YEAR', 'NOC']], drop_first=True)  # 将分类变量进行独热编码
    y = df['gold_count']  # 使用金牌数作为因变量
    # 数据标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    return X_scaled, y

if df_usa.shape[0] > 0:
    X_usa, y_usa = prepare_data(df_usa)
else:
    X_usa, y_usa = None, None

if df_rou.shape[0] > 0:
    X_rou, y_rou = prepare_data(df_rou)
else:#比赛结束前最后一天售后群发布无水印可视化结果+无标注代码【可直接提交】为了防止倒卖， 论文写作过程中遗留数个致命问题，无关代码，该问题解决方式仅在官网授权售后群答疑，盗卖方式购买资料不提供答疑，感谢理解 美赛资料助攻购买链接+说明https://docs.qq.com/doc/p/f3dc6bffbf4dab58dbdfd3e5e5de18a2ad974216
    X_rou, y_rou = None, None

if X_usa is not None and X_rou is not None:
    # 划分训练集和测试集
    X_usa_train, X_usa_test, y_usa_train, y_usa_test = train_test_split(X_usa, y_usa, test_size=0.2, random_state=42)
    X_rou_train, X_rou_test, y_rou_train, y_rou_test = train_test_split(X_rou, y_rou, test_size=0.2, random_state=42)

    # 使用Lasso回归进行建模
    lasso_model_usa = Lasso(alpha=0.1)  # Lasso回归，alpha为正则化参数，可以调节
    lasso_model_usa.fit(X_usa_train, y_usa_train)

    lasso_model_rou = Lasso(alpha=0.1)  # Lasso回归，alpha为正则化参数，可以调节
    lasso_model_rou.fit(X_rou_train, y_rou_train)

    # 预测
    y_usa_pred = lasso_model_usa.predict(X_usa_test)
    y_rou_pred = lasso_model_rou.predict(X_rou_test)

    mse_usa = mean_squared_error(y_usa_test, y_usa_pred)
    mse_rou = mean_squared_error(y_rou_test, y_rou_pred)

    print("USA Model:")
    print("Coefficients:", lasso_model_usa.coef_)
    print("Intercept:", lasso_model_usa.intercept_)

    print("\nRomania Model:")
    print("Coefficients:", lasso_model_rou.coef_)
    print("Intercept:", lasso_model_rou.intercept_)

    print("\nCoach Change Effect Comparison:")
    print(f"USA coach change effect: {lasso_model_usa.coef_[0]}")
    print(f"Romania coach change effect: {lasso_model_rou.coef_[0]}")
else:
    print("Error: No data for either USA or Romania.")
